import os, sys, random, copy
import numpy as np
sys.path.append("..")
import topology_engineer, traffic_engineer, path_selector, traffic_history_reader, traffic_pb2, network_controller
from aurora_network import *

def randomly_swap_logical_topol(nblocks, adj_matrix):
	list1 = []
	list2 = []
	# does at least 10 swaps
	swapped_adj_matrix = copy.deepcopy(adj_matrix)
	for i in range(nblocks):
		for j in range(nblocks):
			if i != j:
				list1.append((i,j))
				list2.append((i,j))
	numswaps = min(max(int(100 * random.random()), 10), len(list1))
	random.shuffle(list1)
	random.shuffle(list2)
	index = 0
	for _ in range(numswaps):
		if list1[index] == list2[index]:
			index += 1
			continue
		else:
			i1 = list1[index][0]
			j1 = list1[index][1]
			i2 = list2[index][0]
			j2 = list2[index][1]
			tmp = swapped_adj_matrix[i1][j1]
			swapped_adj_matrix[i1][j1] = swapped_adj_matrix[i2][j2]
			swapped_adj_matrix[i2][j2] = tmp
		index += 1
	return swapped_adj_matrix

def check_for_empty_flows(aurora_network, traffic_snapshots):
	non_zero_flows = {}
	nblocks = aurora_network.get_num_blocks()
	for snapshot in traffic_snapshots:
		tm = aurora_network.get_traffic_matrix_from_traffic_snapshot(snapshot)
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					if tm[i][j] > 0:
						non_zero_flows[(i,j)] = True
	zero_flows = []
	for i in range(nblocks):
			for j in range(nblocks):
				if i != j and (i,j) not in non_zero_flows:
					zero_flows.append( (i, j, ) )	
	print("zeros flows : {}".format(zero_flows))
	print("fraction of zero flows to total flows: {}".format(float(len(zero_flows))/(nblocks * (nblocks - 1))))
	return zero_flows

# Writes the wcmp weights file
def write_wcmp_weight_file(wcmp_path_weights_filename, nblocks, routing_weights, traffic_matrix):
	str_builder = ""
	#format must be : path len, weight, node1, node2, .... node p, when p = path len
	for path in routing_weights.keys():
		src = path[0]
		dst = path[-1]
		if traffic_matrix[src][dst] == 0:
			continue
		weight = routing_weights[path]
		path_len = len(path)
		if weight >= 0.000001:
			line = "{},{}".format(path_len, weight)
			for node in path:
				line += ",{}".format(node)
			str_builder += (line + "\n")
	with open(wcmp_path_weights_filename, "w+") as f:
		f.write(str_builder)
	return

# Writes the traffic communication probability file
def write_traffic_probability_file(traffic_probabilities_filename, nblocks, traffic_matrix):
	normalized_tm = np.zeros((nblocks, nblocks, ))
	traffic_sum = sum([sum(x) for x in traffic_matrix])
	for i in range(nblocks):
		for j in range(nblocks):
			if traffic_matrix[i][j] > 0:
				normalized_tm[i][j] = traffic_matrix[i][j] / traffic_sum
	str_builder = "#tor_pair_id,src,dst,pdf_num_bytes\n"
	pair = 0
	for i in range(nblocks):
		for j in range(nblocks):
			if i != j and traffic_matrix[i][j] > 0:
				str_builder += ("{},{},{},{}\n".format(pair, i, j, "{:.6E}".format(normalized_tm[i][j])))
				pair += 1
	str_builder += "\n"	
	with open(traffic_probabilities_filename, "w+") as f:
		f.write(str_builder)
	return

## Writes the .topology file 
def write_topology_file(topology_filename, adj_matrix):
	nblocks = len(adj_matrix)
	str_builder = "# aurora with number of pods: {}".format(nblocks)
	str_builder += "\n"
	num_switches = nblocks
	str_builder += ("|V|={}".format(num_switches) + "\n")
	# form a full-mesh
	total_links = sum([sum(x) for x in adj_matrix])
	num_edges = nblocks * (nblocks - 1)
	str_builder += ("|E|={}".format(int(total_links)) + "\n")
	str_builder += ("ToRs=incl_range(" + str(0) + "," + str(num_switches - 1) + ")\n")
	str_builder += ("Servers=incl_range(" + str(0) + "," + str(num_switches - 1) + ")\n")
	str_builder += ("Switches=set()\n\n")
	for src in range(nblocks):
		for dst in range(nblocks):
			lc = 0.
			while lc < adj_matrix[src][dst]:
				str_builder += ("{} {}".format(src, dst) + "\n")
				lc += 1
	with open(topology_filename, "w+") as f:
		f.write(str_builder)
	return

# Writes the simulation parameters into the .properties file
def write_netbench_property_file(config_filename, topology_file, traffic_file, wcmp_weight_filename, output_dir, output_base_dir, load_level):
	# Topology 
	str_builder = "# Topology\n"
	str_builder += ("scenario_topology_file=" + topology_file + "\n") 
	#str_builder += ("scenario_topology_extend_with_servers=regular\n") 
	#str_builder += ("scenario_topology_extend_servers_per_tl_node=" + str(1) + "\n")
	str_builder += "\n"

	# Run info 
	str_builder += "# Run Info\n"
	str_builder += "run_time_s=1\n"
	str_builder += "run_folder_name={}\n".format(output_dir)
	str_builder += "run_folder_base_dir={}\n".format(output_base_dir)
	#str_builder += "analysis_command=python analyze.py\n"
	str_builder += "enable_log_flow_throughput=true\n"
	str_builder += "seed=8278897294\n"
	str_builder += "\n"

	# Network device
	str_builder += "# Network Device\n"
	str_builder += "transport_layer=simple_dctcp\n"
	str_builder += "network_device=wcmp_switch\n"
	str_builder += "network_device_routing=wcmp\n"
	str_builder += "wcmp_path_weights_filename={}\n".format(wcmp_weight_filename)
	str_builder += "network_device_intermediary=identity\n"
	str_builder += "\n"

	# Link & output port
	str_builder += "# Link & output port\n"
	#str_builder += "output_port=ecn_tail_drop\n"
	str_builder += "output_port=ecn_tail_drop_diff_queue_size\n"
	str_builder += "output_port_max_queue_size_bytes=150000\n"
	str_builder += "output_port_ecn_threshold_k_bytes=30000\n"
	str_builder += "link=perfect_simple\n"
	str_builder += "link_delay_ns=600\n"
	str_builder += "link_bandwidth_bit_per_ns={}\n".format(100)
	str_builder += "\n"

	#Traffic
	str_builder += "# Traffic\n"
	str_builder += "traffic=poisson_arrival\n"	
	str_builder += "traffic_lambda_flow_starts_per_s={}\n".format(load_level)
	#str_builder += "traffic_flow_size_dist=pfabric_web_search_upper_bound\n" 
	str_builder += "traffic_flow_size_dist=pfabric_data_mining_upper_bound\n" 
	str_builder += ("traffic_probabilities_file=" + traffic_file +"\n\n")
	with open(config_filename, 'w+') as f:
		f.write(str_builder)
	return


## main executable
def main_diff_mlu(dcn_name, timestamp, block_params, load_level, avehopcount):
	# generate ONE traffic matrix
	# Step 1 : Read Traffic snapshots and also build the topologies
	print("Reading traffic snapshot history from files\nDCN name:{}\n".format(dcn_name))
	block_names_list, traffic_snapshots = traffic_history_reader.TrafficSnapshotReader(dcn_name).get_traffic_snapshots()
	traffic_snapshots_timestamps = [x.time_info.start_time for x in traffic_snapshots]
	timestamp_index = int(0.5 * len(traffic_snapshots_timestamps))
	index = 0
	for ts in traffic_snapshots_timestamps:
		if ts == timestamp:
			timestamp_index = index
			break
		index += 1
	print("Reading traffic files: COMPLETE.......\n")
	
	# Step 2 : Contruct the physical network itself
	aurora_network = AuroraNetwork(dcn_name, block_params, block_names_list)
	nblocks = aurora_network.get_num_blocks()
	zero_flows = check_for_empty_flows(aurora_network, traffic_snapshots)

	traffic_matrix = aurora_network.get_traffic_matrix_from_traffic_snapshot(traffic_snapshots[timestamp_index])
	# Step 4 : Construct the topology engineers and traffic engineers and path selectors
	# 4.1 : Path selector
	all_interblock_paths = path_selector.PathSelector(aurora_network, use_multihop=True).get_all_paths()

	perfect_toe = topology_engineer.IdealTopologyEngineer(aurora_network, all_interblock_paths, traffic_snapshots)
	perfect_te = traffic_engineer.IdealTrafficEngineer(aurora_network, all_interblock_paths)
	ideal_fractional_topology = perfect_toe._topology_engineer(traffic_snapshots[timestamp_index], all_interblock_paths)
	## round?
	ideal_integer_topology = aurora_network.round_fractional_topology_giant_switch(ideal_fractional_topology, zero_flows, rounding_simple=False)
	#aurora_network.print_stats(ideal_fractional_topology)
	#aurora_network.print_stats(ideal_integer_topology)
	perfect_toe_perfect_te_netcontroller = network_controller.NetworkController(aurora_network, perfect_toe, perfect_te, len(traffic_snapshots))
	uniform_routing_weights = perfect_te.compute_path_weights(perfect_toe._generate_uniform_logical_topology(), traffic_matrix)
	integer_routing_weights = perfect_te.compute_path_weights(ideal_integer_topology, traffic_matrix)
	fractional_routing_weights = perfect_te.compute_path_weights(ideal_fractional_topology, traffic_matrix)
	perf_integer = perfect_toe_perfect_te_netcontroller._evaluate_snapshot_performance(timestamp, traffic_matrix, ideal_integer_topology, integer_routing_weights, return_lu_distribution=False)
	perf_fractional = perfect_toe_perfect_te_netcontroller._evaluate_snapshot_performance(timestamp, traffic_matrix, ideal_fractional_topology, fractional_routing_weights, return_lu_distribution=False)
	perf_uniform = perfect_toe_perfect_te_netcontroller._evaluate_snapshot_performance(timestamp, traffic_matrix, perfect_toe._generate_uniform_logical_topology(), uniform_routing_weights, return_lu_distribution=False)
	print("integer perf: {}".format(perf_integer))
	print("uniform perf: {}".format(perf_uniform))
	print("fractional perf: {}".format(perf_fractional))

	topology_list = [ideal_integer_topology,]
	mlu_list = [perf_integer[0],]
	for _ in range(20):
		adj_matrix = randomly_swap_logical_topol(nblocks, ideal_integer_topology)
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j and adj_matrix[i][j] == 0:
					adj_matrix[i][j] = 1.
		topology_list.append(adj_matrix)
		routing_weights = perfect_te.compute_path_weights_fix_avehopcount(adj_matrix, traffic_matrix, avehopcount)
		perf = perfect_toe_perfect_te_netcontroller._evaluate_snapshot_performance(timestamp, traffic_matrix, adj_matrix, routing_weights, return_lu_distribution=False)
		mlu_list.append(perf[0])
	print mlu_list
	# start writing to directories
	os.chdir("./temp/diff_mlu/")
	assert(len(mlu_list) == len(topology_list))
	output_base_dir = "/Users/minyee/src/robust_toe/packet_level_simulation/temp/diff_mlu"
	for index in range(len(topology_list)):
		mlu = mlu_list[index]
		mlu_str = "{:.3f}".format(mlu)
		mlu_str = mlu_str.replace('.', 'p')
		if not os.path.exists("./mlu{}/".format(mlu_str)):
			subdir = "mlu{}/".format(mlu_str)
			os.mkdir(subdir)
			topology_filename = output_base_dir + "/" + subdir + "topol_nblock{}.topology".format(nblocks)
			traffic_filename = output_base_dir + "/" + subdir + "traffic_nblock{}.txt".format(nblocks)
			wcmp_filename = output_base_dir + "/" + subdir + "wcmp_weights_nblock{}.txt".format(nblocks)
			property_filename = output_base_dir + "/" + subdir + "sim_parameters.properties"
			write_topology_file(topology_filename, topology_list[index])
			write_traffic_probability_file(traffic_filename, nblocks, traffic_matrix)
			wcmp_weight = perfect_te.compute_path_weights(topology_list[index], traffic_matrix)
			write_wcmp_weight_file(wcmp_filename, nblocks, wcmp_weight, traffic_matrix)
			write_netbench_property_file(property_filename, topology_filename, traffic_filename, wcmp_filename, "mlu{}".format(mlu_str), output_base_dir, load_level)
	return


def same_traffic_matrix_different_hop_count(dcn_name, timestamp, block_params, load_level):
	# generate ONE traffic matrix
	# Step 1 : Read Traffic snapshots and also build the topologies
	print("Reading traffic snapshot history from files\nDCN name:{}\n".format(dcn_name))
	block_names_list, traffic_snapshots = traffic_history_reader.TrafficSnapshotReader(dcn_name).get_traffic_snapshots()
	traffic_snapshots_timestamps = [x.time_info.start_time for x in traffic_snapshots]
	timestamp_index = int(0.5 * len(traffic_snapshots_timestamps))
	index = 0
	avehopcounts = np.arange(1.1, 2.0, 0.1)
	for ts in traffic_snapshots_timestamps:
		if ts == timestamp:
			timestamp_index = index
			break
		index += 1
	print("Reading traffic files: COMPLETE.......\n")
	
	# Step 2 : Contruct the physical network itself
	aurora_network = AuroraNetwork(dcn_name, block_params, block_names_list)
	nblocks = aurora_network.get_num_blocks()
	#zero_flows = check_for_empty_flows(aurora_network, traffic_snapshots)

	traffic_matrix = aurora_network.get_traffic_matrix_from_traffic_snapshot(traffic_snapshots[timestamp_index])
	# Step 4 : Construct the topology engineers and traffic engineers and path selectors
	# 4.1 : Path selector
	all_interblock_paths = path_selector.PathSelector(aurora_network, use_multihop=True).get_all_paths()

	perfect_toe = topology_engineer.IdealTopologyEngineer(aurora_network, all_interblock_paths, traffic_snapshots)
	perfect_te = traffic_engineer.IdealTrafficEngineer(aurora_network, all_interblock_paths)
	#ideal_fractional_topology = perfect_toe._topology_engineer(traffic_snapshots[timestamp_index], all_interblock_paths)
	uniform_fractional_topology = perfect_toe._generate_uniform_logical_topology()
	uniform_fractional_topology = aurora_network.round_fractional_topology_giant_switch(uniform_fractional_topology, [])
	print uniform_fractional_topology
	
	perfect_toe_perfect_te_netcontroller = network_controller.NetworkController(aurora_network, perfect_toe, perfect_te, len(traffic_snapshots))
	


	### Uniform Topology
	if not os.path.exists("./temp/diff_ahc/load_{}_uniform".format(load_level)):
		os.mkdir("./temp/diff_ahc/load_{}_uniform".format(load_level))
	os.chdir("./temp/diff_ahc/load_{}_uniform".format(load_level))
	output_base_dir = "/Users/minyee/src/robust_toe/packet_level_simulation/temp/diff_ahc/load_{}_uniform".format(load_level)

	topology_filename = output_base_dir + "/" + "topol_nblock{}.topology".format(nblocks)
	traffic_filename = output_base_dir + "/" + "traffic_nblock{}.txt".format(nblocks)
	write_topology_file(topology_filename, uniform_fractional_topology)
	write_traffic_probability_file(traffic_filename, nblocks, traffic_matrix)
	for ahc in avehopcounts:
		
		ahc_str = "{:.1f}".format(ahc)
		ahc_str = ahc_str.replace('.', 'p')
		subdir = "ahc{}/".format(ahc_str)
		if not os.path.exists("./{}".format(subdir)):
			os.mkdir(subdir)

		wcmp_filename = output_base_dir + "/" + subdir + "wcmp_weights_nblock{}.txt".format(nblocks)
		property_filename = output_base_dir + "/" + subdir + "sim_parameters.properties"
		routing_weights = perfect_te.compute_path_weights_fix_avehopcount(uniform_fractional_topology, traffic_matrix, ahc)
		write_wcmp_weight_file(wcmp_filename, nblocks, routing_weights, traffic_matrix)
		write_netbench_property_file(property_filename, topology_filename, traffic_filename, wcmp_filename, "ahc{}".format(ahc_str), output_base_dir, load_level)

	### ToE topology
	os.chdir("../../..")
	if not os.path.exists("./temp/diff_ahc/load_{}_toe".format(load_level)):
		os.mkdir("./temp/diff_ahc/load_{}_toe".format(load_level))
	os.chdir("./temp/diff_ahc/load_{}_toe".format(load_level))
	output_base_dir = "/Users/minyee/src/robust_toe/packet_level_simulation/temp/diff_ahc/load_{}_toe".format(load_level)
	toe_fractional_topology = perfect_toe._topology_engineer(traffic_snapshots[timestamp_index], all_interblock_paths)
	toe_fractional_topology = aurora_network.round_fractional_topology_giant_switch(toe_fractional_topology, [])
	topology_filename = output_base_dir + "/" + "topol_nblock{}.topology".format(nblocks)
	traffic_filename = output_base_dir + "/" + "traffic_nblock{}.txt".format(nblocks)
	write_topology_file(topology_filename, toe_fractional_topology)
	write_traffic_probability_file(traffic_filename, nblocks, traffic_matrix)
	for ahc in avehopcounts:
		
		ahc_str = "{:.1f}".format(ahc)
		ahc_str = ahc_str.replace('.', 'p')
		subdir = "ahc{}/".format(ahc_str)
		if not os.path.exists("./{}".format(subdir)):
			os.mkdir(subdir)

		wcmp_filename = output_base_dir + "/" + subdir + "wcmp_weights_nblock{}.txt".format(nblocks)
		property_filename = output_base_dir + "/" + subdir + "sim_parameters.properties"
		routing_weights = perfect_te.compute_path_weights_fix_avehopcount(uniform_fractional_topology, traffic_matrix, ahc)
		write_wcmp_weight_file(wcmp_filename, nblocks, routing_weights, traffic_matrix)
		write_netbench_property_file(property_filename, topology_filename, traffic_filename, wcmp_filename, "ahc{}".format(ahc_str), output_base_dir, load_level)
	#for ahc in avehopcounts:
	#	routing_weights = perfect_te.compute_path_weights_fix_avehopcount(uniform_fractional_topology, traffic_matrix, ahc)
	#	perf_tuple = perfect_toe_perfect_te_netcontroller._evaluate_snapshot_performance(timestamp, traffic_matrix, uniform_fractional_topology, routing_weights, return_lu_distribution=False)
	#	ahc = perf_tuple[3]
	#	mlu = perf_tuple[0]
	#	lu50 = perf_tuple[2]
	#	lu50_simulated.append(lu50)
	#	ahc_simulated.append(ahc)
	#	mlu_simulated.append(mlu)
	#print ahc_simulated
	#print mlu_simulated
	#print lu50_simulated
	return

# arrival rate
# total_traffic(gbps) = arrival_rate * ave_flow size

if __name__ == "__main__":
	dcn_name = "ju1_dls08"
	timestamp = 1526540400
	superblock_num_links = 256
	border_router_num_links = 256
	superblock_link_capacity = 100
	border_router_link_capacity = 100

	block_params = {BlockType.SUPERBLOCK : {}, BlockType.BORDER_ROUTER : {}}
	block_params[BlockType.SUPERBLOCK]["link capacity"] = float(superblock_link_capacity)
	block_params[BlockType.SUPERBLOCK]["num links"] = float(superblock_num_links)
	block_params[BlockType.BORDER_ROUTER]["link capacity"] = float(border_router_link_capacity) # in gbps
	block_params[BlockType.BORDER_ROUTER]["num links"] = float(border_router_num_links)
	load_level = 100000
	#same_traffic_matrix_different_hop_count(dcn_name, timestamp, block_params, load_level)
	avehopcount = 1.1
	main_diff_mlu(dcn_name, timestamp, block_params, load_level, avehopcount)